-- | Multivariate polynomials in the power basis

module Ply (C(..), T, toTree, dRecT, lo) where

import qualified Tree
import qualified Iv
import Interlude

class Functor p => C p where
  const :: a -> p a
  isConst :: p a -> Maybe a
  var :: Num a => Int -> Iv.T a -> p a
  vars :: p a -> Int
  eval :: Num a => p a -> [a] -> a

  compose :: Num (p a) => {-Uni-}p a -> p a -> p a
  compose g f = eval (fmap Ply.const g) [f]

  scale :: (Num a, Num (p a)) => p a -> [Iv.T a] -> p a
  scale p = eval (fmap Ply.const p) . zipWith var [0..]

newtype T a = P { toTree :: Tree.T a } deriving (Eq,Show,Functor,Arbitrary)
-- Invariant: no empty branches (Tree.Br [])

instance C T where
  const = P . Tree.Lf

  isConst (P (Tree.Lf c)) = Just c
  isConst _               = Nothing

  var = (P.) . varT where
    varT 0 x = Tree.Br [Tree.Lf (Iv.lo x), Tree.Lf (Iv.width x)]
    varT i x = Tree.Br [varT (i-1) x]

  vars = Tree.depth . toTree

  eval = hornerT . toTree where
    hornerT (Tree.Lf c) _      = c
    hornerT (Tree.Br b) (x:xs) = horner0 (map (`hornerT` xs) b) x
    horner0 []     _ = 0
    horner0 [p]    _ = p
    horner0 (p:ps) x = p + x * horner0 ps x

instance Num a => Num (T a) where
  P x + P y = P (x + y)
  (*) = \(P x) (P y) -> P (x |*| y) where
    Tree.Lf x |*| Tree.Lf y = Tree.Lf $! x*y
    Tree.Lf x |*| Tree.Br y = Tree.Br $! map' (fmap (x*)) y
    Tree.Br x |*| Tree.Lf y = Tree.Br $! map' (fmap (y*)) x
    Tree.Br x |*| Tree.Br y = Tree.Br $! x|**|y
    []      |**| _          = []
    (x0:x') |**| y@(y0:y')  = (x0|*|y0) !: zipKeep (+) (map' (x0|*|) y') (x'|**|y)

  negate = fmap negate
  fromInteger = P . Tree.Lf . fromInteger
  signum = error "signum[Ply.T]"
  abs = error "abs[Ply.T]"

instance Fractional a => Fractional (T a) where
  recip (P (Tree.Lf a)) = P $ Tree.Lf $ recip a
  fromRational = P . Tree.Lf . fromRational

-- | Differential recursion. Operates on raw trees (@Tree.t@).
dRecT :: Num t => a -> (Tree.T t -> a -> a) -> Tree.T t -> a
dRecT z f = \(Tree.Br p) -> dr p where
  dr []     = z
  dr (p0:p) = f p0 $ dr $ zipWith (fmap . (*) . fromInteger) [1..] p

-- | [lo k p] gives the polynomial composed of the monomials in which
-- all exponents are at most [k]. Its coefficients form a hyper-cube.
lo :: Int -> T a -> T a
lo k = P . loT . toTree where
  loT (Tree.Lf c) = Tree.Lf c
  loT (Tree.Br b) = Tree.Br $ map loT $ take k b

{-# SPECIALIZE instance Num (T Double) #-}
{-# SPECIALIZE instance Fractional (T Double) #-}
